import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs

def loadDataSet(fileName):      
    """导入数据

    Args:
        fileName (string): 导入文档的名称

    Returns:
        numpy.ndarray: 原始数据集，以二维ndarry数组形式存储（每个点的坐标）
    """
    dataMat = []                
    fr = open(fileName)
    for line in fr.readlines():  #逐行读取文档
        curLine = line.strip().split('\t')  #文档中是以tab作分隔符
        fltLine = list(map(float,curLine)) 
        dataMat.append(fltLine)
    return np.array(dataMat)

def distEclud(vecA, vecB):
    """计算两个向量之间的距离（欧氏距离）

    Args:
        vecA (numpy.ndarray): 向量A
        vecB (numpy.ndarray): 向量B

    Returns:
        float: 向量间的距离
    """
    return np.sqrt(sum(np.power(vecA - vecB, 2))) #la.norm(vecA-vecB)

def randCent(dataSet, k):
    """随机化初始质心

    Args:
        dataSet (numpy.ndarray): 原始数据集，以二维ndarry数组形式存储
        k (int): 初始质心的个数

    Returns:
        numpy.matrix: 初始质心的坐标
    """
    n = np.shape(dataSet)[1]
    centroids = np.mat(np.zeros((k, n))) #创建k个初始质心
    for j in range(n):
        minJ = min(dataSet[:, j])
        rangeJ = float(max(dataSet[:, j]) - minJ)
        centroids[:, j] = np.mat(minJ + rangeJ * np.random.rand(k, 1))
    return centroids
    
def kMeans(dataSet, k, distMeas=distEclud, createCent=randCent):
    """kMeans算法

    Args:
        dataSet (numpy.ndarray): 原始数据集，以二维ndarry数组形式存储
        k (int): 划分的簇的数目
        distMeas (function, optional): 生成初始质心的函数. Defaults to distEclud.
        createCent (function, optional): 计算欧几里得距离的函数. Defaults to randCent.

    Returns:
        numpy.ndarray, numpy.matrix: 质心的坐标, 所有点分配给指定质心的信息（所属的簇，与质心的距离）
    """
    m = np.shape(dataSet)[0]    #m为数据点的数目
    clusterAssment = np.mat(np.zeros((m,2)))#记录所有点的指派信息（指派给了哪个质心，距离是多少）
    centroids = np.array(createCent(dataSet, k))  #选择K个点作为初始质心
    clusterChanged = True   
    while clusterChanged:
        clusterChanged = False
        for i in range(m):  #循环中，将每一个点指派到最近的质心，i表示第i个点
            minIndex = -1; minDist = np.inf #记录与哪个质心最近，距离是多少
            for j in range(k):  #循环中，将第i个点与每一个质心进行比较
                distJI = distMeas(dataSet[i,:], centroids[j,:])  #计算第j个质心和第i个点的距离（i是外循环）
                if distJI < minDist:
                    minDist = distJI; minIndex = j
            if clusterAssment[i,0] != minIndex: #判断这个被指派的质心是否发生了变化
                clusterChanged = True   #只要有一个点被指派的质心发生了变化，while循环会一直进行下去
            clusterAssment[i,:] = minIndex,minDist**2
        for cent in range(k):   #重新计算每个簇的质心
            ptsInClust = dataSet[np.nonzero(clusterAssment[:,0].A==cent)[0]] 
            #ptsInClust记录分配给第cent个簇的所有点（坐标），np.nonzero接收的是布尔类型数组，返回值的是索引的数组
            centroids[cent,:] = np.mean(ptsInClust, axis=0) 
            #根据这些点的两个坐标重新计算每个簇的质心，np.mean(axis=0)会计算每一列上的均值 
    return centroids, clusterAssment

def createDataSet():
    X, y = make_blobs(n_samples=500,n_features=2,centers=4,random_state=10)
    # X是包含两个特征的数据集，y是真实标签，在聚类时并不使用
    return X, y

def scatter(dataSet):
    """绘制散点图（二维坐标系）

    Args:
        dataSet (numpy.ndarray): 原始数据集，以二维ndarry数组形式存储
    """
    X = np.array(dataSet)   # 将原始数据集转化为ndarray数组，方便切片
    plt.figure(figsize=(10,7), dpi=200)
    plt.scatter(X[:,0],X[:,1],marker='o',s=50)  #这里可以自己尝试输出X[:,0]和X[:,1]，看看是什么
    plt.savefig("./fig1.png")
    #plt.show()

def scatterSep(dataSet, clusterAssment, centroids):
    """根据最终聚类结果绘制散点图（用不同的颜色代表不同的簇）

    Args:
        dataSet (numpy.ndarray)): 原始数据集，以二维ndarry数组形式存储
        clusterAssment (numpy.matrix): 每个点的最终指派信息
        centroids (numpy.ndarray): 质心的坐标
    """
    y = np.array([int(i[0]) for i in np.array(clusterAssment)])
    X = np.array(dataSet)
    color = ["red","black","green","blue"]
    #这里的color只填入了4个颜色，最多给四个簇染色，如果kMeans选择4个簇以上，则color需要添加其他颜色
    plt.figure(figsize=(10,7), dpi=200)
    k = 4   #簇的数目
    for i in range(k):
        plt.scatter(X[y==i,0],X[y==i,1],marker='o',s=50,c=color[i]) #maker设置点的形状，s设置点的大小
    plt.scatter(centroids[:,0], centroids[:,1], marker='*',s=100,c="orange")
    plt.savefig("./fig2.png")
    #plt.show()

if __name__ == '__main__':
    dataSet = loadDataSet('testSet.txt')
    # dataSet, y_real = createDataSet()
    centroids, clusterAssment = kMeans(dataSet, 4)
    scatterSep(dataSet, clusterAssment, centroids)
    print(np.mean(np.asarray(clusterAssment), axis=0)[1])   #打印SSE